"""
MetalX - Agent

"""

import socket
import threading
import json
import time
import uuid
import asyncio

from db import util


panel_log_db = util.JsonDatabase("panel_log.json")


def insert_log(level: str, data: str):
    panel_log_db.insert("log", {"level": level, "data": data, "timestamp": time.time()})


class Controller:
    def __init__(self, host="0.0.0.0", port=65432):
        """
        初始化API服务器

        Args:
            host (str, optional): 监听的主机地址. 默认值为 '0.0.0.0'.
            port (int, optional): 监听的端口号. 默认值为 65432.
        """
        self.host = host
        self.port = port
        self.server_socket = None
        self.clients = {}  # 存储客户端连接信息 {client_id: (conn, addr)}
        self.running = False
        self.lock = threading.Lock()  # 用于线程安全操作
        self.messageQueue = {}  # {"request_id":message["request_id"],"data":message}

    def start(self):
        """启动服务器"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(10)
        self.on_log(f"Controller启动，监听 {self.host}:{self.port}")

        self.running = True
        accept_thread = threading.Thread(target=self.accept_clients, daemon=True)
        accept_thread.start()

    def accept_clients(self):
        """
        接受客户端连接
        函数会一直运行直到手动停止或因异常结束
        """
        while self.running:
            try:
                conn, addr = self.server_socket.accept()
                client_id = str(uuid.uuid4())
                with self.lock:
                    self.clients[client_id] = (conn, addr)

                self.on_connect(f"新的客户端连接: ID={client_id}, 地址={addr}")
                client_thread = threading.Thread(
                    target=self.handle_client, args=(client_id, conn, addr), daemon=True
                )
                client_thread.start()
            except Exception as e:
                if self.running:
                    self.on_error(f"接受客户端连接时出错: {e}")

    def handle_client(self, client_id, conn, addr):
        """
        处理与客户端的通信
        函数会一直运行直到手动停止或因异常结束
        """
        try:
            while self.running:
                data = conn.recv(4096)
                if not data:
                    break

                try:
                    message = json.loads(data.decode("utf-8"))
                    self.process_client_message(client_id, message)
                except json.JSONDecodeError:
                    self.on_error(f"来自客户端 {client_id} 的无效JSON数据: {data}")
        except (ConnectionResetError, ConnectionAbortedError):
            self.on_disconnect(f"客户端 {client_id} 断开连接")
        except Exception as e:
            self.on_error(f"处理客户端 {client_id} 时出错: {e}")
        finally:
            self.remove_client(client_id)

    def process_client_message(self, client_id, message: dict):
        """处理来自客户端的消息"""
        if "type" not in message:
            self.on_error(f"来自客户端 {client_id} 的无效消息: 缺少消息类型")
            return

        
        self.on_message(f"来自客户端 {client_id} 消息: {message}")
        if message.get("request_id"):
            self.messageQueue.update({message["request_id"]: message})
            self.on_log(f"当前消息队列：{self.messageQueue}")

    async def send_command(self, client_id, command, args=None):
        """
        向指定客户端发送命令，并等待与 request_id 相符的结果
        :param client_id: 客户端ID
        :param command: 命令名称
        :param args: 命令参数
        :return: 命令执行结果，如果超时则返回 None
        """
        request_id =  str(uuid.uuid4())
        with self.lock:
            if client_id not in self.clients:
                self.on_error(f"错误: 客户端ID {client_id} 不存在")
                return None

            conn, _ = self.clients[client_id]

        message = {
            "type": "command",
            "command": command,
            "args": args or None,
            "request_id": request_id,
        }

        try:
            conn.sendall(json.dumps(message).encode("utf-8"))
        except Exception as e:
            self.on_error(f"向客户端 {client_id} 发送命令失败: {e}")
            self.remove_client(client_id)
            return None

        # 设置超时时间（秒）
        timeout = 10
        start_time = time.time()

        while True:
            # 检查是否超时
            if time.time() - start_time > timeout:
                self.on_error(f"请求 {request_id} 超时，未收到响应")
                return None

            with self.lock:
                if self.messageQueue.get(request_id):
                    result = self.messageQueue.get(request_id)
                    del self.messageQueue[request_id]
                    return result

            # 短暂休眠，避免CPU占用过高
            await asyncio.sleep(0.2)

    async def broadcast_command(self, command, args=None):
        """
        向所有客户端广播命令
        :param command: 命令名称
        :param args: 命令参数
        :return: 成功发送的客户端数量
        """
        result = {}
        
        with self.lock:
            client_ids = list(self.clients.keys())

        for client_id in client_ids:
            with self.lock:
                if client_id not in self.clients:
                    self.on_error(f"错误: 客户端ID {client_id} 不存在")
                    continue

                conn, _ = self.clients[client_id]

            message = {
                "type": "command",
                "command": command,
                "args": args or None,
                "request_id": client_id, #使用client_id，可以复用，方便计数
            }

            try:
                conn.sendall(json.dumps(message).encode("utf-8"))
            except Exception as e:
                self.on_error(f"向客户端 {client_id} 发送命令失败: {e}")
                self.remove_client(client_id)
                continue
        
        # 请求发完了开始找结果，设置超时时间（秒）
        timeout = 10
        start_time = time.time()

        while True:
            # 检查是否超时
            if time.time() - start_time > timeout:
                self.on_error(f"请求超时，未收到响应")
                break

            with self.lock:
                for request_id in client_ids:
                    if self.messageQueue.get(request_id):
                        result.update({request_id:self.messageQueue.get(request_id)}) 
                        del self.messageQueue[request_id]
                    return result

            # 短暂休眠，避免CPU占用过高
            await asyncio.sleep(0.2)
        
    
    

    def list_clients(self):
        """
        获取当前连接的客户端列表
        :return: 客户端ID和地址的列表
        """
        with self.lock:
            return [
                {"id": cid, "address": f"{addr[0]}:{addr[1]}"}
                for cid, (_, addr) in self.clients.items()
            ]

    def remove_client(self, client_id):
        """移除客户端"""
        with self.lock:
            if client_id in self.clients:
                conn, _ = self.clients[client_id]
                try:
                    conn.close()
                except:
                    pass
                del self.clients[client_id]
                self.on_disconnect(f"客户端 {client_id} 已断开连接")

    def stop(self):
        """停止服务器"""
        # 先标记服务器停止运行，让接受客户端连接的线程退出循环
        self.running = False

        try:
            # 关闭服务器套接字，解除 accept 方法的阻塞
            if self.server_socket:
                self.server_socket.close()
        except Exception as e:
            self.on_error(f"关闭服务器套接字时出错: {e}")

        # 关闭所有客户端连接
        with self.lock:
            for client_id in list(self.clients.keys()):
                try:
                    conn, _ = self.clients[client_id]
                    conn.close()
                except Exception as e:
                    self.on_error(f"关闭客户端 {client_id} 连接时出错: {e}")
                finally:
                    del self.clients[client_id]
                    self.on_disconnect(f"客户端 {client_id} 已断开连接")

        self.on_log("服务器已停止")

    # region 回调函数
    def on_connect(self, info: str):
        print(f"[Controller]连接成功：{info}")
        insert_log("log", f"[Controller]连接成功：{info}")

    def on_disconnect(self, info: str):
        print(f"[Controller]断开连接: {info}")
        insert_log("log", f"[Controller]断开连接: {info}")

    def on_message(self, message: str):
        print(f"[Controller]收到消息: {message}")
        insert_log("log", f"[Controller]收到消息: {message}")

    def on_error(self, error: str):
        print(f"[Controller]发生错误: {error}")
        insert_log("error", f"[Controller]发生错误: {error}")

    def on_log(self, message: str):
        print(f"[Controller]事件: {message}")
        insert_log("log", f"[Controller]事件: {message}")

    # endregion 回调函数


if __name__ == "__main__":
    server = Controller()
    server.start()

    # 保持主线程运行
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        server.stop()
